#include <iostream>
#include <sstream>
#include <cstring>
#include <vector>
#include <queue>
#include <algorithm>
#include <random>
#include "getopt.h"
#include <sys/stat.h>
#include "config-fnt.h"
#include "graph.priv.h"
#include "stream.h"
#include "public/utils.h"

/* Options */
static const char* opts=":o:s:d:l:D:S:hv";
static const struct option opts_long[]=
{
	{"outdir", 1, nullptr, 'o'},
	{"size", 1, nullptr, 's'},
	{"distance", 1, nullptr, 'd'},
	{"length", 1, nullptr, 'l'},
	{"safe-dist", 1, nullptr, 'D'},
	{"soma-dist", 1, nullptr, 'S'},
	{"help", 0, nullptr, 'h'},
	{"version", 0, nullptr, 'v'},
	{nullptr, 0, nullptr, 0}
};

/* Print version information and exit */
void version(const char* name) {
	std::cout<<name<<" (" PACKAGE_NAME_LONG ") " PACKAGE_VERSION "\n";
	std::cout<<"Copyright (C) 2015-2016 " PACKAGE_ORG "\n";
	std::cout<<"License GPLv3+: GNU GPL version 3 or later <http://gnu.org/licenses/gpl.html>\n\n";
	std::cout<<"Written by: " PACKAGE_AUTHORS "\n";
	exit(0);
}
// //
/* Print usage information and exit */
void usage(const char* name, int ecode) {
	if(ecode) {
		std::cerr<<"Try '"<<name<<" --help' for more information.\n";
		exit(ecode);
	}
	std::cout<<"Usage: "<<name<<" [OPTION]... FNT_FILE\n";
	std::cout<<
		"Split a FNT file and find errors.\n\n"
		"Options:\n"
		"  -o, --outdir=DIR  [OPTIONAL] Save splitted files to DIR. To current\n"
		"                       working directory by default.\n\n"
		"  -s, --size=SIZE   [OPTIONAL] Desired size of splitted files (in\n"
		"                       number of branches). 256 by default.\n\n"
		"  -d, --distance=D  [OPTIONAL] Set the distance for branch collision\n"
		"                       test (in um). 0.6um by default.\n"
		"  -l, --length=L    [OPTIONAL] Set the length for branch collision\n"
		"                       test (in um). 10um by default.\n\n"
		"  -D, --safe-dist=D [OPTIONAL] Set the conservative distance for\n"
		"                       collision tests with edges that are not part\n"
		"                       of a neuron. 2um by default.\n\n"
		"  -S, --soma-dist=S [OPTIONAL] Set the distance for soma collision\n"
		"                       test (in um). 20um by default.\n\n"
		"  -h, --help        [OPTIONAL] Print this help, and do nothing else.\n"
		"  -v, --version     [OPTIONAL] Print version information and exit.\n\n"
		"Report bugs to <" PACKAGE_BUGREPORT ">\n"
		PACKAGE_NAME_LONG " home page: <" PACKAGE_URL ">\n";
	exit(ecode);
}

/* Ensure the dir exist. Try create if not. Return true on success. */
bool ensureDir(const char* dir) {
	struct stat buf;
	int r=stat(dir, &buf);
	if(r!=-1 && S_ISDIR(buf.st_mode)) {
		return true;
	}
#ifdef BUILD_MINGW
	r=mkdir(dir);
#else
	r=mkdir(dir, 0777);
#endif
	if(r!=-1)
		return true;
	std::cerr<<"Cannot create dir '"<<dir<<"': "<<strerror(errno)<<".\n";
	return false;
}

Graph load_graph(const char* filename, std::vector<std::string>& header) {
	Graph gr=GraphPriv::alloc();
	try {
		InputStream fs;
		if(!fs.open(filename))
			throwError("Failed to open file.");
		fs.precision(13);
		fs.setf(std::ios::scientific);

		std::string line;
		std::getline(fs, line);
		if(!fs)
			throwError("Failed to read magic header");
		bool gzipped=false;
		if(line==FNTZ_MAGIC) {
			gzipped=true;
		} else if(line!=FNT_MAGIC) {
			throwError("Wrong magic: ", line);
		}
		header.push_back(line);

		std::getline(fs, line);
		if(!fs)
			throwError("Failed to read URL line");
		while(line!="BEGIN_TRACING_DATA") {
			header.push_back(line);
			std::getline(fs, line);
			if(!fs)
				throwError("Failed to read line");
		}

		if(gzipped) {
			if(!fs.pushGzip())
				throwError("Failed to start decompressing");
		}
		auto gp=GraphPriv::get(gr);
		gp->load(fs);
		if(gzipped) {
			if(!fs.pop())
				throwError("Failed to finish decompressing");
		}
		if(!fs.close())
			throwError("Failed to close file.");
		gp->updateModel();
	} catch(const std::exception& e) {
		GraphPriv::free(gr);
	}
	return gr;
}
bool save_graph(Graph gr, const std::string& filename, const std::vector<std::string>& header) {
	bool ret=true;
	try {
		OutputStream fs;
		if(!fs.open(filename))
			throwError("Failed to open file");
		fs.precision(13);
		fs.setf(std::ios::scientific);

		bool gzipped=false;
		if(header[0]==FNTZ_MAGIC)
			gzipped=true;

		const std::vector<std::string> stringsToSkip{
			"core.current.",
				"core.target."
		};
		for(auto& l: header) {
			bool skip=false;
			for(auto& s: stringsToSkip) {
				if(l.compare(0, s.length(), s)==0) {
					skip=true;
					break;
				}
			}
			if(skip)
				continue;
			fs<<l<<'\n';
			if(!fs)
				throwError("Failed to write line");
		}

		fs<<"BEGIN_TRACING_DATA"<<'\n';
		if(!fs)
			throwError("Failed to write data start mark");

		if(gzipped) {
			if(!fs.pushGzip())
				throwError("Failed to start compressing");
		}
		GraphPriv::get(gr)->save(fs);
		if(gzipped) {
			if(!fs.pop())
				throwError("Failed to finish compressing");
		}
		if(!fs.close())
			throwError("Failed to close file");
	} catch(const std::exception& e) {
		ret=false;
	}
	return ret;
}

enum ADJ_REASON {
	NO_ADJ=0,
	ADJ_CONN=1<<0,
	ADJ_SOMA_COLL=1<<1,
	ADJ_EDGE_COLL=1<<2,
};

struct AdjItem {
	size_t i;
	int r;
};

bool edg_connected(Edge e1, Edge e2) {
	if(e1.leftVertex()==e2.leftVertex())
		return true;
	if(e1.leftVertex()==e2.rightVertex())
		return true;
	if(e1.rightVertex()==e2.rightVertex())
		return true;
	if(e1.rightVertex()==e2.leftVertex())
		return true;
	return false;
}
bool soma_collision(Edge e1, Edge e2, double somadist) {
	auto n1=e1.tree();
	if(!n1)
		return false;
	auto n2=e2.tree();
	if(!n2)
		return false;
	if(n1==n2)
		return false;
	if(n1.root().point().distTo(n2.root().point())<somadist)
		return true;
	return false;
}
inline double calcDist(const Point& p0, const Point& p1, const Point& p2) {
	double p0x=p0.x();
	double p0y=p0.y();
	double p0z=p0.z();
	double p1x=p1.x();
	double p1y=p1.y();
	double p1z=p1.z();
	double p2x=p2.x();
	double p2y=p2.y();
	double p2z=p2.z();
	double ux=p1x-p2x;
	double uy=p1y-p2y;
	double uz=p1z-p2z;
	double vx=p2x-p0x;
	double vy=p2y-p0y;
	double vz=p2z-p0z;
	double a=ux*ux+uy*uy+uz*uz;
	double b=ux*vx+uy*vy+uz*vz;
	double c=vx*vx+vy*vy+vz*vz;
	double t=-b/a;
	if(t<0) t=0;
	if(t>1) t=1;
	return sqrt(a*t*t+2*b*t+c);
}
bool edge_collision(Edge e1, Edge e2, double edgdist, double edglen,
		double mingap1, double maxgap1, double mingap2, double maxgap2) {
	auto& pts1=e1.points();
	auto& pts2=e2.points();
	auto n1=pts1.size();
	auto n2=pts2.size();
	for(size_t i=0; i<n2;) {
		double min_dist=INFINITY;
		for(size_t j=0; j+1<n1;) {
			auto dist=calcDist(pts2[i], pts1[j], pts1[j+1]);
			if(dist<min_dist)
				min_dist=dist;
			if(dist<edgdist) {
				auto j0=j;
				double len=0;
				auto i1=i+1;
				while(i1<n2 && len<edglen && dist<edgdist) {
					dist=INFINITY;
					auto l=pts2[i1].distTo(pts2[i1-1]);
					size_t w=(l+edgdist)/mingap1+1;
					if(w>100) w=100;
					size_t j0b=j0;
					size_t j0e=j0;
					if(j0>=w)
						j0b=j0-w;
					else
						j0b=0;
					if(j0+w+1<n1)
						j0e=j0+w+1;
					else
						j0e=n1-1;
					for(size_t j1=j0b; j1<j0e; j1++) {
						auto d=calcDist(pts2[i1], pts1[j1], pts1[j1+1]);
						if(d<dist) {
							dist=d;
							j0=j1;
						}
					}
					if(dist<edgdist) {
						len+=l;
					}
					i1++;
				}
				if(len>=edglen)
					return true;
			}
			int inc=(dist-edgdist)/maxgap1+1;
			if(inc<1) inc=1;
			j+=inc;
		}
		int inc=(min_dist-edgdist)/maxgap2;
		if(inc<1) inc=1;
		i+=inc;
	}
	return false;
}
bool edge_collision_safe(Edge e1, Edge e2, double edgdist,
		double mingap1, double maxgap1, double mingap2, double maxgap2) {
	auto& pts1=e1.points();
	auto& pts2=e2.points();
	auto n1=pts1.size();
	auto n2=pts2.size();
	bool e1t=e1.tree()?true:false;
	bool e2t=e2.tree()?true:false;
	for(size_t i=0; i<n2;) {
		double min_dist=INFINITY;
		for(size_t j=0; j+1<n1;) {
			auto dist=calcDist(pts2[i], pts1[j], pts1[j+1]);
			if(dist<min_dist)
				min_dist=dist;
			if(dist<edgdist) {
				return true;
			}
			int inc=(dist-edgdist)/maxgap1+1;
			if(e1t) {
				if(inc<1) inc=1;
			} else {
				inc=n1-2;
				if(inc<1) inc=1;
			}
			j+=inc;
		}
		int inc=(min_dist-edgdist)/maxgap2;
		if(e2t) {
			if(inc<1) inc=1;
		} else {
			inc=n2-1;
			if(inc<1) inc=1;
		}
		i+=inc;
	}
	return false;
}
std::vector<std::vector<AdjItem>> calc_adj(Graph gr, double edgdist, double edglen, double safedist, double somadist) {
	auto n=gr.edges().size();
	std::vector<double> mingaps(n);
	std::vector<double> maxgaps(n);
	for(size_t i=0; i<n; i++) {
		double mingap=INFINITY;
		double maxgap=-INFINITY;
		auto& pts=gr.edges()[i].points();
		for(size_t j=0; j+1<pts.size(); j++) {
			auto d=pts[j].distTo(pts[j+1]);
			if(d<mingap) mingap=d;
			if(d>maxgap) maxgap=d;
		}
		mingaps[i]=mingap;
		maxgaps[i]=maxgap;
	}
	std::vector<std::vector<AdjItem>> adjs(n);
//#pragma omp parallel for
	for(size_t i=0; i<n; i++) {
		for(size_t j=0; j<i; j++) {
			Edge ei=gr.edges()[i];
			Edge ej=gr.edges()[j];
			int r=NO_ADJ;
			if(edg_connected(ei, ej))
				r|=ADJ_CONN;
			if(soma_collision(ei, ej, somadist))
				r|=ADJ_SOMA_COLL;
			if(ei.tree() && ej.tree()) {
				if(edge_collision(ei, ej, edgdist, edglen, mingaps[i], maxgaps[i], mingaps[j], maxgaps[j])) {
					r|=ADJ_EDGE_COLL;
					auto ep=EdgePriv::get(ei);
					ep->type=1;
					ep=EdgePriv::get(ej);
					ep->type=1;
				}
			} else {
				if(edge_collision_safe(ei, ej, safedist, mingaps[i], maxgaps[i], mingaps[j], maxgaps[j])) {
					r|=ADJ_EDGE_COLL;
					auto ep=EdgePriv::get(ei);
					ep->type=1;
					ep=EdgePriv::get(ej);
					ep->type=1;
				}
			}
			if(r!=NO_ADJ) {
				adjs[i].emplace_back(AdjItem{j, r});
				adjs[j].emplace_back(AdjItem{i, r});
			}
		}
//#pragma omp critical
		std::cerr<<"finished "<<i<<'\n';
	}
	return adjs;
}

enum ERROR_TYPE {
	WITHOUT_ERROR=0,
	JUNK=1<<0,//
	UNMARKED=1<<1,//
	MISMARKED=1<<2,//
	UNFINISHED=1<<3,//
	SOMA_COLLISION=1<<4,//
	EDGE_COLLISION=1<<5,//
};
struct Component {
	int err;
	std::vector<size_t> edges;
	std::vector<size_t> verts;
	std::vector<size_t> nrns;
};
std::vector<Component> find_components(Graph gr, const std::vector<std::vector<AdjItem>>& adjs) {
	std::vector<Component> comps;
	auto n=gr.edges().size();
	std::vector<bool> edg_dirty(n, false);
	std::vector<bool> vert_dirty(gr.vertices().size(), false);
	std::vector<bool> nrn_dirty(gr.trees().size(), false);
	while(true) {
		size_t seed=SIZE_MAX;
		for(size_t i=0; i<n; i++) {
			if(!edg_dirty[i]) {
				seed=i;
				break;
			}
		}
		if(seed==SIZE_MAX)
			break;

		Component comp{WITHOUT_ERROR, {}, {}, {}};
		std::queue<size_t> que{};
		que.push(seed);
		while(!que.empty()) {
			auto i=que.front();
			que.pop();
			if(!edg_dirty[i]) {
				comp.edges.emplace_back(i);
				for(auto& adj: adjs[i]) {
					if(adj.r&ADJ_EDGE_COLL)
						comp.err|=EDGE_COLLISION;
					if(adj.r&ADJ_SOMA_COLL)
						comp.err|=SOMA_COLLISION;
					que.push(adj.i);
				}
				edg_dirty[i]=true;
			}
		}
		for(auto i: comp.edges) {
			auto e=gr.edges()[i];
			auto j=e.leftVertex().index();
			auto k=e.rightVertex().index();
			if(!vert_dirty[j]) {
				comp.verts.push_back(j);
				vert_dirty[j]=true;
			}
			if(!vert_dirty[k]) {
				comp.verts.push_back(k);
				vert_dirty[k]=true;
			}
			if(!e.leftVertex().finished() || !e.rightVertex().finished()) {
				comp.err|=UNFINISHED;
			}
		}
		for(auto i: comp.edges) {
			auto nrn=gr.edges()[i].tree();
			if(nrn) {
				auto j=nrn.index();
				if(!nrn_dirty[j]) {
					comp.nrns.push_back(j);
					nrn_dirty[j]=true;
				}
			} else {
				comp.err|=JUNK;
			}
		}
		comps.push_back(std::move(comp));
	}
	for(auto v: edg_dirty) {
		if(!v)
			throwError("Untouched edge");
	}
	for(auto v: vert_dirty) {
		if(!v)
			throwError("Untouched vertex");
	}
	for(auto v: nrn_dirty) {
		if(!v)
			throwError("Untouched neuron");
	}
	return comps;
}

int main(int argc, char* argv[]) {
	const char* outdir=".";

	size_t outsize=256;
	double edgdist=0.6;
	double edglen=10;
	double safedist=2;
	double somadist=20;

	// Parse options
	int opt;
	int prev_optind=optind;
	while((opt=getopt_long(argc, argv, opts, opts_long, nullptr))!=-1) {
		switch(opt) {
			case '?':
				std::cerr<<argv[0]<<": unrecognized option '";
				std::cerr<<argv[prev_optind]<<"'\n";
				usage(argv[0], -1);
			case ':':
				std::cerr<<argv[0]<<": option '"<<argv[prev_optind];
				std::cerr<<"' requires an argument\n";
				usage(argv[0], -1);
			case 'h':
				usage(argv[0], 0);
			case 'v':
				version(argv[0]);
			case 'o':
				outdir=optarg;
				break;
			case 's':
				outsize=atol(optarg);
				break;
			case 'd':
				edgdist=atof(optarg);
				break;
			case 'l':
				edglen=atof(optarg);
				break;
			case 'D':
				safedist=atof(optarg);
				break;
			case 'S':
				somadist=atof(optarg);
				break;
			default:
				std::cerr<<"Unexpected error.\n";
				usage(argv[0], -1);
		}
		prev_optind=optind;
	}

	if(optind+1<argc) {
		std::cerr<<"Too many input FNT files.\n";
		usage(argv[0], -1);
	}
	if(optind>=argc) {
		std::cerr<<"No input FNT file specified.\n";
		usage(argv[0], -1);
	}

	// Ensure arguments are correct
	if(!ensureDir(outdir)) {
		return -1;
	}

	{
		if(outsize<=0) {
			std::cerr<<"SIZE should be positive.\n";
			return -1;
		}
		if(edgdist<0) {
			std::cerr<<"D should be positive.\n";
			return -1;
		}
		if(edglen<=0) {
			std::cerr<<"L should be positive.\n";
			return -1;
		}
		if(safedist<0) {
			std::cerr<<"Conservative distance should be positive.\n";
			return -1;
		}
		if(somadist<0) {
			std::cerr<<"S should be positive.\n";
			return -1;
		}
	}

	std::vector<std::string> header;
	auto graph=load_graph(argv[optind], header);
	if(!graph) {
		std::cerr<<"Failed to load FNT file.\n";
		return -1;
	}

	auto adjs=calc_adj(graph, edgdist, edglen, safedist, somadist);
	auto comps=find_components(graph, adjs);
	std::cerr<<"comps: "<<comps.size()<<"\n";

	std::vector<Component> comps_err;
	std::vector<Component> comps_noerr;
	for(auto& comp: comps) {
		if(comp.err==WITHOUT_ERROR) {
			comps_noerr.push_back(std::move(comp));
		} else {
			comps_err.push_back(std::move(comp));
		}
	}

	size_t idx;
	for(idx=0; idx<comps_err.size(); idx++) {
		auto grout=GraphPriv::alloc();
		auto gp=GraphPriv::get(grout);
		auto& comp=comps_err[idx];
		for(auto i: comp.edges) {
			gp->edges.push_back(graph.edges()[i]);
		}
		for(auto i: comp.verts) {
			gp->vertices.push_back(graph.vertices()[i]);
		}
		for(auto i: comp.nrns) {
			gp->trees.push_back(graph.trees()[i]);
		}
		gp->updateModel();
		std::ostringstream oss;
		oss<<outdir<<'/';
		oss.width(4);
		oss.fill('0');
		oss<<idx;
		if(comp.err&JUNK) oss<<"_JUNK";
		if(comp.err&UNMARKED) oss<<"_UNMARKED";
		if(comp.err&MISMARKED) oss<<"_MISMARKED";
		if(comp.err&UNFINISHED) oss<<"_UNFINISHED";
		if(comp.err&SOMA_COLLISION) oss<<"_SOMA-COLLISION";
		if(comp.err&EDGE_COLLISION) oss<<"_EDGE-COLLISION";
		oss<<".fnt";
		auto fn=oss.str();
		bool ret=save_graph(grout, fn, header);
		if(!ret) {
			std::cerr<<"Failed to save: "<<fn<<"\n";
		}
		GraphPriv::free(grout);
	}
	std::vector<size_t> shuffled(comps_noerr.size());
	for(size_t i=0; i<shuffled.size(); i++) {
		shuffled[i]=i;
	}
	std::shuffle(shuffled.begin(), shuffled.end(), std::default_random_engine{});
	{
		auto grout=GraphPriv::alloc();
		auto gp=GraphPriv::get(grout);
		size_t nerri=0;
		while(nerri<comps_noerr.size()) {
			auto& comp=comps_noerr[shuffled[nerri]];
			for(auto i: comp.edges) {
				gp->edges.push_back(graph.edges()[i]);
			}
			for(auto i: comp.verts) {
				gp->vertices.push_back(graph.vertices()[i]);
			}
			for(auto i: comp.nrns) {
				gp->trees.push_back(graph.trees()[i]);
			}

			if(grout.edges().size()>=outsize || nerri+1==comps_noerr.size()) {
				gp->updateModel();
				std::ostringstream oss;
				oss<<outdir<<'/';
				oss.width(4);
				oss.fill('0');
				oss<<idx;
				oss<<".fnt";
				auto fn=oss.str();
				bool ret=save_graph(grout, fn, header);
				if(!ret) {
					std::cerr<<"Failed to save: "<<fn<<"\n";
				}
				idx++;
				gp->vertices.clear();
				gp->edges.clear();
				gp->trees.clear();
			}
			nerri++;
		}
		GraphPriv::free(grout);
	}
	return 0;
}

